#include "srgan.h"
#include <random>
#include "utils/logging.h"

static std::vector<std::string> g_classes = {};

SRGAN::SRGAN()
{
    engine_ = CreateRKNNEngine();
    input_tensor_.data = nullptr;
    want_float_ = false; // 是否使用浮点数版本的后处理
    ready_ = false;
}

SRGAN::~SRGAN()
{
    // release input tensor and output tensor
    NN_LOG_DEBUG("release input tensor");
    if (input_tensor_.data != nullptr)
    {
        free(input_tensor_.data);
        input_tensor_.data = nullptr;
    }
    NN_LOG_DEBUG("release output tensor");
    for (auto &tensor : output_tensors_)
    {
        if (tensor.data != nullptr)
        {
            free(tensor.data);
            tensor.data = nullptr;
        }
    }
}

nn_error_e SRGAN::LoadModel(const char *model_path)
{
    auto ret = engine_->LoadModelFile(model_path);
    if (ret != NN_SUCCESS)
    {
        NN_LOG_ERROR("yolov8 load model file failed");
        return ret;
    }
    // get input tensor
    auto input_shapes = engine_->GetInputShapes();

    // check number of input and n_dims
    if (input_shapes.size() != 1)
    {
        NN_LOG_ERROR("input tensor number is not 1, but %ld", input_shapes.size());
        return NN_RKNN_INPUT_ATTR_ERROR;
    }
    nn_tensor_attr_to_cvimg_input_data(input_shapes[0], input_tensor_);
    input_tensor_.data = malloc(input_tensor_.attr.size);

    auto output_shapes = engine_->GetOutputShapes();
    if (output_shapes.size() != 1)
    {
        NN_LOG_ERROR("output tensor number is not 1, but %ld", output_shapes.size());
        return NN_RKNN_OUTPUT_ATTR_ERROR;
    }
    if (output_shapes[0].type == NN_TENSOR_FLOAT16)
    {
        want_float_ = true;
        NN_LOG_WARNING("output tensor type is float16, want type set to float32");
    }
    for (int i = 0; i < output_shapes.size(); i++)
    {
        tensor_data_s tensor;
        tensor.attr.n_elems = output_shapes[i].n_elems;
        tensor.attr.n_dims = output_shapes[i].n_dims;
        for (int j = 0; j < output_shapes[i].n_dims; j++)
        {
            tensor.attr.dims[j] = output_shapes[i].dims[j];
        }
        // output tensor needs to be float32
        // tensor.attr.type = want_float_ ? NN_TENSOR_FLOAT : output_shapes[i].type;
        tensor.attr.type = want_float_ ? NN_TENSOR_FLOAT : NN_TENSOR_INT8;
        tensor.attr.index = 0;
        tensor.attr.size = output_shapes[i].n_elems * nn_tensor_type_to_size(tensor.attr.type);
        tensor.data = malloc(tensor.attr.size);
        output_tensors_.push_back(tensor);
        out_zps_.push_back(output_shapes[i].zp);
        out_scales_.push_back(output_shapes[i].scale);
    }
 
    ready_ = true;
    return NN_SUCCESS;
}


nn_error_e SRGAN::Preprocess(const cv::Mat &img)
{

    // 比例
    cv::Mat img_resized;
    cv::resize(img, img_resized, cv::Size(input_tensor_.attr.dims[2], input_tensor_.attr.dims[1]), 0, 0, cv::INTER_LINEAR);
    memcpy(input_tensor_.data, img_resized.data, input_tensor_.attr.size);
    return NN_SUCCESS;
}

nn_error_e SRGAN::Inference()
{
    std::vector<tensor_data_s> inputs;
    inputs.push_back(input_tensor_);
    return engine_->Run(inputs, output_tensors_, want_float_);
}

nn_error_e SRGAN::Postprocess( cv::Mat &result_img)
{

    // 获取输出数据指针（假设输出为float32）
    float* output_data = (float*)output_tensors_[0].data;

    // 获取维度信息
    int batch = output_tensors_[0].attr.dims[0];  // 通常是1
    int channels = output_tensors_[0].attr.dims[1]; // 应该是3
    int height = output_tensors_[0].attr.dims[2];
    int width = output_tensors_[0].attr.dims[3];

    // 创建OpenCV矩阵（NCHW -> HWC）
    cv::Mat output_mat(height, width, CV_32FC3);
    for (int h = 0; h < height; h++) {
        for (int w = 0; w < width; w++) {
            for (int c = 0; c < 3; c++) {
                float value = output_data[c * height * width + h * width + w];
                // value = (value + 1.0f) * 127.5f;
                output_mat.at<cv::Vec3f>(h, w)[c] = value;
            }
        }
    }

    // 转换数值范围到[0,255]
    cv::normalize(output_mat, output_mat, 0, 255, cv::NORM_MINMAX);
    output_mat.convertTo(result_img, CV_8UC3);

    // 颜色空间转换（如果模型输出是RGB）
    cv::cvtColor(result_img, result_img, cv::COLOR_RGB2BGR);
    
    return NN_SUCCESS;
}



nn_error_e SRGAN::Run(const cv::Mat &img, cv::Mat &result_img)
{

    // letterbox后的图像
    Preprocess(img);
    // 推理
    Inference();
    // 后处理
    Postprocess(result_img);

    return NN_SUCCESS;
}
